# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-08-04 23:32
# File     : websocket_routes.py
# Project  : codebuddy_craft
# Desc     : WebSocket接口

# backend/api/websocket_routes.py
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, Query
from fastapi.websockets import WebSocketState
from sqlalchemy.ext.asyncio import AsyncSession
from typing import Dict, List, Set, Optional, Any
import json
import asyncio
import uuid
from datetime import datetime
from enum import Enum

from backend.database.connection import get_db
from backend.core.dependencies import get_user_from_token, get_current_user
from backend.services.project_service import ProjectService
from backend.services.file_service import FileService
from backend.services.llm_service import LLMService
from backend.models.user import User
from backend.utils.logger import get_logger

logger = get_logger(__name__)
from backend.utils.rate_limiter import RateLimiter

router = APIRouter()

# WebSocket 消息类型
class MessageType(str, Enum):
    # 连接管理
    CONNECT = "connect"
    DISCONNECT = "disconnect"
    PING = "ping"
    PONG = "pong"
    
    # 聊天相关
    CHAT_MESSAGE = "chat_message"
    CHAT_RESPONSE = "chat_response"
    CHAT_TYPING = "chat_typing"
    CHAT_HISTORY = "chat_history"
    
    # 代码协作
    CODE_CHANGE = "code_change"
    CODE_CURSOR = "code_cursor"
    CODE_SELECTION = "code_selection"
    CODE_SAVE = "code_save"
    
    # 项目管理
    PROJECT_UPDATE = "project_update"
    FILE_CREATE = "file_create"
    FILE_UPDATE = "file_update"
    FILE_DELETE = "file_delete"
    
    # 系统通知
    NOTIFICATION = "notification"
    ERROR = "error"
    SUCCESS = "success"


class WebSocketConnection:
    """WebSocket 连接封装"""
    
    def __init__(
        self, 
        websocket: WebSocket, 
        user_id: str, 
        connection_id: str,
        project_id: Optional[str] = None
    ):
        self.websocket = websocket
        self.user_id = user_id
        self.connection_id = connection_id
        self.project_id = project_id
        self.connected_at = datetime.utcnow()
        self.last_ping = datetime.utcnow()
        self.subscribed_channels: Set[str] = set()
    
    async def send_message(self, message_type: MessageType, data: Dict[str, Any]):
        """发送消息"""
        try:
            if self.websocket.client_state == WebSocketState.CONNECTED:
                message = {
                    "type": message_type.value,
                    "data": data,
                    "timestamp": datetime.utcnow().isoformat(),
                    "connection_id": self.connection_id
                }
                await self.websocket.send_text(json.dumps(message, ensure_ascii=False))
        except Exception as e:
            logger.error(f"发送WebSocket消息失败: {str(e)}")
    
    async def send_error(self, error_message: str, error_code: str = "GENERAL_ERROR"):
        """发送错误消息"""
        await self.send_message(MessageType.ERROR, {
            "message": error_message,
            "code": error_code
        })
    
    async def send_success(self, message: str, data: Optional[Dict[str, Any]] = None):
        """发送成功消息"""
        await self.send_message(MessageType.SUCCESS, {
            "message": message,
            "data": data or {}
        })


class WebSocketManager:
    """WebSocket 连接管理器"""
    
    def __init__(self):
        # 所有活跃连接
        self.connections: Dict[str, WebSocketConnection] = {}
        
        # 用户连接映射
        self.user_connections: Dict[str, Set[str]] = {}
        
        # 项目房间映射
        self.project_rooms: Dict[str, Set[str]] = {}
        
        # 频道订阅
        self.channel_subscriptions: Dict[str, Set[str]] = {}
        
        # 服务依赖
        self.project_service = ProjectService()
        self.file_service = FileService()
        self.llm_service = LLMService()
        self.rate_limiter = RateLimiter()
        
        # 启动心跳检测任务
        asyncio.create_task(self._heartbeat_task())
    
    async def connect(
        self, 
        websocket: WebSocket, 
        user_id: str, 
        project_id: Optional[str] = None
    ) -> WebSocketConnection:
        """建立WebSocket连接"""
        
        connection_id = str(uuid.uuid4())
        connection = WebSocketConnection(websocket, user_id, connection_id, project_id)
        
        # 保存连接
        self.connections[connection_id] = connection
        
        # 添加到用户连接映射
        if user_id not in self.user_connections:
            self.user_connections[user_id] = set()
        self.user_connections[user_id].add(connection_id)
        
        # 加入项目房间
        if project_id:
            await self._join_project_room(connection_id, project_id)
        
        logger.info(f"WebSocket连接建立 - 用户: {user_id}, 连接ID: {connection_id}")
        
        # 发送连接成功消息
        await connection.send_success("连接成功", {
            "connection_id": connection_id,
            "user_id": user_id,
            "project_id": project_id
        })
        
        return connection
    
    async def disconnect(self, connection_id: str):
        """断开WebSocket连接"""
        
        if connection_id in self.connections:
            connection = self.connections[connection_id]
            
            # 从项目房间移除
            if connection.project_id:
                await self._leave_project_room(connection_id, connection.project_id)
            
            # 从频道订阅中移除
            for channel in connection.subscribed_channels:
                self._unsubscribe_channel(connection_id, channel)
            
            # 从用户连接映射中移除
            user_connections = self.user_connections.get(connection.user_id, set())
            user_connections.discard(connection_id)
            if not user_connections:
                del self.user_connections[connection.user_id]
            
            # 移除连接
            del self.connections[connection_id]
            
            logger.info(f"WebSocket连接断开 - 用户: {connection.user_id}, 连接ID: {connection_id}")
    
    async def send_to_user(self, user_id: str, message_type: MessageType, data: Dict[str, Any]):
        """向特定用户发送消息"""
        
        user_connections = self.user_connections.get(user_id, set())
        for connection_id in user_connections.copy():
            connection = self.connections.get(connection_id)
            if connection:
                await connection.send_message(message_type, data)
    
    async def send_to_project(self, project_id: str, message_type: MessageType, data: Dict[str, Any], exclude_user: Optional[str] = None):
        """向项目房间发送消息"""
        
        room_connections = self.project_rooms.get(project_id, set())
        for connection_id in room_connections.copy():
            connection = self.connections.get(connection_id)
            if connection and (not exclude_user or connection.user_id != exclude_user):
                await connection.send_message(message_type, data)
    
    async def send_to_channel(self, channel: str, message_type: MessageType, data: Dict[str, Any]):
        """向频道发送消息"""
        
        channel_connections = self.channel_subscriptions.get(channel, set())
        for connection_id in channel_connections.copy():
            connection = self.connections.get(connection_id)
            if connection:
                await connection.send_message(message_type, data)
    
    async def broadcast(self, message_type: MessageType, data: Dict[str, Any]):
        """广播消息给所有连接"""
        
        for connection in self.connections.values():
            await connection.send_message(message_type, data)
    
    async def _join_project_room(self, connection_id: str, project_id: str):
        """加入项目房间"""
        
        if project_id not in self.project_rooms:
            self.project_rooms[project_id] = set()
        self.project_rooms[project_id].add(connection_id)
        
        # 通知其他用户有新用户加入
        connection = self.connections.get(connection_id)
        if connection:
            await self.send_to_project(
                project_id, 
                MessageType.PROJECT_UPDATE,
                {
                    "action": "user_joined",
                    "user_id": connection.user_id,
                    "project_id": project_id
                },
                exclude_user=connection.user_id
            )
    
    async def _leave_project_room(self, connection_id: str, project_id: str):
        """离开项目房间"""
        
        if project_id in self.project_rooms:
            self.project_rooms[project_id].discard(connection_id)
            if not self.project_rooms[project_id]:
                del self.project_rooms[project_id]
        
        # 通知其他用户有用户离开
        connection = self.connections.get(connection_id)
        if connection:
            await self.send_to_project(
                project_id,
                MessageType.PROJECT_UPDATE,
                {
                    "action": "user_left",
                    "user_id": connection.user_id,
                    "project_id": project_id
                },
                exclude_user=connection.user_id
            )
    
    def _subscribe_channel(self, connection_id: str, channel: str):
        """订阅频道"""
        
        if channel not in self.channel_subscriptions:
            self.channel_subscriptions[channel] = set()
        self.channel_subscriptions[channel].add(connection_id)
        
        connection = self.connections.get(connection_id)
        if connection:
            connection.subscribed_channels.add(channel)
    
    def _unsubscribe_channel(self, connection_id: str, channel: str):
        """取消订阅频道"""
        
        if channel in self.channel_subscriptions:
            self.channel_subscriptions[channel].discard(connection_id)
            if not self.channel_subscriptions[channel]:
                del self.channel_subscriptions[channel]
        
        connection = self.connections.get(connection_id)
        if connection:
            connection.subscribed_channels.discard(channel)
    
    async def _heartbeat_task(self):
        """心跳检测任务"""
        
        while True:
            try:
                await asyncio.sleep(30)  # 每30秒检查一次
                
                current_time = datetime.utcnow()
                dead_connections = []
                
                for connection_id, connection in self.connections.items():
                    # 检查超过90秒没有ping的连接
                    if (current_time - connection.last_ping).total_seconds() > 90:
                        dead_connections.append(connection_id)
                    else:
                        # 发送ping
                        try:
                            await connection.send_message(MessageType.PING, {})
                        except Exception:
                            dead_connections.append(connection_id)
                
                # 清理死连接
                for connection_id in dead_connections:
                    await self.disconnect(connection_id)
                
            except Exception as e:
                logger.error(f"心跳检测任务错误: {str(e)}")
    
    def get_project_users(self, project_id: str) -> List[str]:
        """获取项目中的在线用户"""
        
        room_connections = self.project_rooms.get(project_id, set())
        users = set()
        
        for connection_id in room_connections:
            connection = self.connections.get(connection_id)
            if connection:
                users.add(connection.user_id)
        
        return list(users)
    
    def get_connection_stats(self) -> Dict[str, Any]:
        """获取连接统计信息"""
        
        return {
            "total_connections": len(self.connections),
            "total_users": len(self.user_connections),
            "total_projects": len(self.project_rooms),
            "total_channels": len(self.channel_subscriptions)
        }


# 全局WebSocket管理器
websocket_manager = WebSocketManager()


# WebSocket 路由

@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str = Query(..., description="认证token"),
    project_id: Optional[str] = Query(None, description="项目ID"),
    db: AsyncSession = Depends(get_db)
):
    """主WebSocket端点"""
    
    try:
        # 验证用户身份
        user = await get_user_from_token(token, db)
        if not user:
            await websocket.close(code=4001, reason="认证失败")
            return
        
        # 验证项目权限
        if project_id:
            project = await websocket_manager.project_service.get_project(project_id, db)
            if not project or project.user_id != user.id:
                await websocket.close(code=4003, reason="项目权限不足")
                return
        
        # 接受连接
        await websocket.accept()
        
        # 创建连接
        connection = await websocket_manager.connect(websocket, user.id, project_id)
        
        try:
            # 消息处理循环
            while True:
                data = await websocket.receive_text()
                await _handle_websocket_message(connection, data, db)
        
        except WebSocketDisconnect:
            pass
        except Exception as e:
            logger.error(f"WebSocket消息处理错误: {str(e)}")
            await connection.send_error(f"消息处理错误: {str(e)}")
    
    except Exception as e:
        logger.error(f"WebSocket连接错误: {str(e)}")
        try:
            await websocket.close(code=4000, reason=str(e))
        except:
            pass
    
    finally:
        # 清理连接
        if 'connection' in locals():
            await websocket_manager.disconnect(connection.connection_id)


async def _handle_websocket_message(
    connection: WebSocketConnection, 
    raw_data: str, 
    db: AsyncSession
):
    """处理WebSocket消息"""
    
    try:
        # 解析消息
        message = json.loads(raw_data)
        message_type = message.get("type")
        data = message.get("data", {})
        
        # 速率限制检查
        if not await websocket_manager.rate_limiter.check_rate_limit(
            f"ws_{connection.user_id}",
            max_requests=100,
            window_seconds=60
        ):
            await connection.send_error("消息发送太频繁", "RATE_LIMIT_EXCEEDED")
            return
        
        # 更新最后ping时间
        if message_type == MessageType.PONG.value:
            connection.last_ping = datetime.utcnow()
            return
        
        # 路由消息到处理器
        handler = _get_message_handler(message_type)
        if handler:
            await handler(connection, data, db)
        else:
            await connection.send_error(f"未知消息类型: {message_type}", "UNKNOWN_MESSAGE_TYPE")
    
    except json.JSONDecodeError:
        await connection.send_error("消息格式错误", "INVALID_JSON")
    except Exception as e:
        logger.error(f"处理WebSocket消息失败: {str(e)}")
        await connection.send_error(f"消息处理失败: {str(e)}", "PROCESSING_ERROR")


def _get_message_handler(message_type: str):
    """获取消息处理器"""
    
    handlers = {
        MessageType.PING.value: _handle_ping,
        MessageType.CHAT_MESSAGE.value: _handle_chat_message,
        MessageType.CHAT_TYPING.value: _handle_chat_typing,
        MessageType.CODE_CHANGE.value: _handle_code_change,
        MessageType.CODE_CURSOR.value: _handle_code_cursor,
        MessageType.CODE_SAVE.value: _handle_code_save,
        MessageType.PROJECT_UPDATE.value: _handle_project_update,
    }
    
    return handlers.get(message_type)


# 消息处理器

async def _handle_ping(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理ping消息"""
    
    connection.last_ping = datetime.utcnow()
    await connection.send_message(MessageType.PONG, {"timestamp": datetime.utcnow().isoformat()})


async def _handle_chat_message(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理聊天消息"""
    
    try:
        message = data.get("message", "")
        project_id = data.get("project_id") or connection.project_id
        
        if not message:
            await connection.send_error("消息内容不能为空", "EMPTY_MESSAGE")
            return
        
        if len(message) > 10000:
            await connection.send_error("消息长度过长", "MESSAGE_TOO_LONG")
            return
        
        # 广播用户消息到项目房间
        if project_id:
            await websocket_manager.send_to_project(
                project_id,
                MessageType.CHAT_MESSAGE,
                {
                    "user_id": connection.user_id,
                    "message": message,
                    "timestamp": datetime.utcnow().isoformat()
                },
                exclude_user=connection.user_id
            )
        
        # 生成AI回复（异步）
        asyncio.create_task(_generate_ai_response(connection, message, project_id, db))
        
    except Exception as e:
        logger.error(f"处理聊天消息失败: {str(e)}")
        await connection.send_error("聊天消息处理失败", "CHAT_ERROR")


async def _handle_chat_typing(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理打字状态"""
    
    project_id = data.get("project_id") or connection.project_id
    is_typing = data.get("is_typing", False)
    
    if project_id:
        await websocket_manager.send_to_project(
            project_id,
            MessageType.CHAT_TYPING,
            {
                "user_id": connection.user_id,
                "is_typing": is_typing,
                "timestamp": datetime.utcnow().isoformat()
            },
            exclude_user=connection.user_id
        )


async def _handle_code_change(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理代码变更"""
    
    try:
        file_id = data.get("file_id")
        changes = data.get("changes", [])
        project_id = data.get("project_id") or connection.project_id
        
        if not file_id:
            await connection.send_error("文件ID不能为空", "MISSING_FILE_ID")
            return
        
        # 验证文件权限
        file = await websocket_manager.file_service.get_file(file_id, db)
        if not file or file.project.user_id != connection.user_id:
            await connection.send_error("文件不存在或无权限", "FILE_ACCESS_DENIED")
            return
        
        # 广播代码变更到项目房间
        if project_id:
            await websocket_manager.send_to_project(
                project_id,
                MessageType.CODE_CHANGE,
                {
                    "file_id": file_id,
                    "user_id": connection.user_id,
                    "changes": changes,
                    "timestamp": datetime.utcnow().isoformat()
                },
                exclude_user=connection.user_id
            )
        
    except Exception as e:
        logger.error(f"处理代码变更失败: {str(e)}")
        await connection.send_error("代码变更处理失败", "CODE_CHANGE_ERROR")


async def _handle_code_cursor(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理光标位置"""
    
    file_id = data.get("file_id")
    cursor_position = data.get("cursor_position", {})
    project_id = data.get("project_id") or connection.project_id
    
    if project_id and file_id:
        await websocket_manager.send_to_project(
            project_id,
            MessageType.CODE_CURSOR,
            {
                "file_id": file_id,
                "user_id": connection.user_id,
                "cursor_position": cursor_position,
                "timestamp": datetime.utcnow().isoformat()
            },
            exclude_user=connection.user_id
        )


async def _handle_code_save(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理代码保存"""
    
    try:
        file_id = data.get("file_id")
        content = data.get("content", "")
        project_id = data.get("project_id") or connection.project_id
        
        if not file_id:
            await connection.send_error("文件ID不能为空", "MISSING_FILE_ID")
            return
        
        # 保存文件
        await websocket_manager.file_service.update_file_content(file_id, content, db)
        
        # 通知保存成功
        await connection.send_success("文件保存成功", {"file_id": file_id})
        
        # 广播文件更新到项目房间
        if project_id:
            await websocket_manager.send_to_project(
                project_id,
                MessageType.FILE_UPDATE,
                {
                    "file_id": file_id,
                    "user_id": connection.user_id,
                    "action": "saved",
                    "timestamp": datetime.utcnow().isoformat()
                },
                exclude_user=connection.user_id
            )
        
    except Exception as e:
        logger.error(f"处理代码保存失败: {str(e)}")
        await connection.send_error("代码保存失败", "CODE_SAVE_ERROR")


async def _handle_project_update(connection: WebSocketConnection, data: Dict[str, Any], db: AsyncSession):
    """处理项目更新"""
    
    project_id = data.get("project_id") or connection.project_id
    action = data.get("action")
    
    if project_id and action:
        await websocket_manager.send_to_project(
            project_id,
            MessageType.PROJECT_UPDATE,
            {
                "user_id": connection.user_id,
                "action": action,
                "data": data.get("update_data", {}),
                "timestamp": datetime.utcnow().isoformat()
            },
            exclude_user=connection.user_id
        )


async def _generate_ai_response(
    connection: WebSocketConnection, 
    message: str, 
    project_id: Optional[str], 
    db: AsyncSession
):
    """生成AI回复"""
    
    try:
        # 构建上下文
        context = {}
        
        if project_id:
            project = await websocket_manager.project_service.get_project(project_id, db)
            if project:
                context["project"] = {
                    "id": project.id,
                    "name": project.name,
                    "description": project.description,
                    "tech_stack": project.tech_stack
                }
        
        # 生成AI回复
        response = await websocket_manager.llm_service.generate_chat_response(
            message=message,
            context=context,
            conversation_history=[]
        )
        
        # 发送AI回复
        await connection.send_message(MessageType.CHAT_RESPONSE, {
            "message": response["content"],
            "type": response.get("type", "text"),
            "metadata": response.get("metadata", {}),
            "timestamp": datetime.utcnow().isoformat()
        })
        
        # 如果项目存在，也广播给其他用户
        if project_id:
            await websocket_manager.send_to_project(
                project_id,
                MessageType.CHAT_RESPONSE,
                {
                    "message": response["content"],
                    "type": response.get("type", "text"),
                    "metadata": response.get("metadata", {}),
                    "from_ai": True,
                    "timestamp": datetime.utcnow().isoformat()
                },
                exclude_user=connection.user_id
            )
    
    except Exception as e:
        logger.error(f"生成AI回复失败: {str(e)}")
        await connection.send_error("AI回复生成失败", "AI_RESPONSE_ERROR")


# 管理端点

@router.get("/ws/stats")
async def get_websocket_stats():
    """获取WebSocket连接统计"""
    
    return websocket_manager.get_connection_stats()


@router.get("/ws/project/{project_id}/users")
async def get_project_online_users(project_id: str):
    """获取项目在线用户"""
    
    users = websocket_manager.get_project_users(project_id)
    return {"project_id": project_id, "online_users": users}


@router.post("/ws/broadcast")
async def broadcast_message(
    message_data: Dict[str, Any],
    current_user: User = Depends(get_current_user)
):
    """广播消息（管理员功能）"""
    
    # 这里可以添加管理员权限检查
    message_type = MessageType(message_data.get("type", MessageType.NOTIFICATION.value))
    data = message_data.get("data", {})
    
    await websocket_manager.broadcast(message_type, data)
    
    return {"message": "消息已广播"}


@router.post("/ws/project/{project_id}/notify")
async def notify_project_users(
    project_id: str,
    notification_data: Dict[str, Any],
    current_user: User = Depends(get_current_user)
):
    """通知项目用户"""
    
    await websocket_manager.send_to_project(
        project_id,
        MessageType.NOTIFICATION,
        notification_data
    )
    
    return {"message": f"已通知项目 {project_id} 的所有用户"}


# 工具函数

def get_websocket_manager() -> WebSocketManager:
    """获取WebSocket管理器实例"""
    return websocket_manager


async def notify_file_change(
    project_id: str, 
    file_id: str, 
    action: str, 
    user_id: str,
    file_data: Optional[Dict[str, Any]] = None
):
    """通知文件变更"""
    
    await websocket_manager.send_to_project(
        project_id,
        MessageType.FILE_UPDATE,
        {
            "file_id": file_id,
            "action": action,
            "user_id": user_id,
            "file_data": file_data,
            "timestamp": datetime.utcnow().isoformat()
        }
    )


async def notify_project_change(
    project_id: str,
    action: str,
    user_id: str,
    project_data: Optional[Dict[str, Any]] = None
):
    """通知项目变更"""
    
    await websocket_manager.send_to_project(
        project_id,
        MessageType.PROJECT_UPDATE,
        {
            "action": action,
            "user_id": user_id,
            "project_data": project_data,
            "timestamp": datetime.utcnow().isoformat()
        }
    )


async def send_system_notification(
    user_id: str,
    title: str,
    message: str,
    notification_type: str = "info"
):
    """发送系统通知给特定用户"""
    
    await websocket_manager.send_to_user(
        user_id,
        MessageType.NOTIFICATION,
        {
            "title": title,
            "message": message,
            "type": notification_type,
            "timestamp": datetime.utcnow().isoformat()
        }
    )
